from mindspore import Tensor
import mindspore
import math
import numpy as np
import mindspore.nn as nn
from model.module_test import Net
from mindspore import context


# class Net(nn.Cell):
#     def __init__(self):
#         super(Net, self).__init__()
#
#     def construct(self, x):
#         return mindspore.Tensor(x, dtype=mindspore.int32)


if __name__ == '__main__':
    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
    b = 3
    crops1 = Tensor(np.ones((b, 4, 3, 224, 224)), mindspore.float32)
    crops2 = Tensor(np.ones((b, 10, 3, 96, 96)), mindspore.float32)
    crops3 = Tensor([2], mindspore.int16)
    text = Tensor(np.ones((b, 77)), mindspore.int64)
    net = Net()
    output = net(crops1, crops2, text)
    print(output)
